import matplotlib.pyplot as plt
import matplotlib
from typing import Tuple, Optional

import cv2
import torch
import numpy as np

from algorithms.SAM.sam import Sam
from algorithms.transforms import ResizeLongestSide


class SamPredictor:
    """
    SAM模型的预测器，集成了预测机制
    """

    def __init__(
            self,
            model: Sam,
            device: torch.device,
            timeout: float = 0.01,
    ) -> None:
        """
        Args:
            model (Sam): 在外部构建好的SAM模型
        """
        super().__init__()
        self.input_size = None
        self.original_size = None
        self.input_w = None
        self.input_h = None
        self.orig_w = None
        self.orig_h = None
        self.features = None
        self.is_image_set = None
        self.model = model
        self.transform = ResizeLongestSide(model.image_encoder.img_size)
        self.reset_image()
        self._task = None
        self.device = device

    def predict(self, data):
        image_cv2 = cv2.imread(data["imagePath"])
        image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
        ori_image_height, ori_image_width, _ = image_cv2.shape
        ratio_height = ori_image_height / data["frontImageHeight"]
        ratio_width = ori_image_width / data["frontImageWidth"]
        self.set_image(image_cv2)
        point = data["points"][0]
        label = data["labels"][0]
        imageConfigX = data["imageConfigX"]
        imageConfigY = data["imageConfigY"]
        input_point = np.array([[point["x"] * ratio_width, point["y"] * ratio_height]])
        input_label = np.array([label[0]])
        masks, scores, logits = self._predict(point_coords=input_point,
                                              point_labels=input_label,
                                              multi_mask_output=True)
        batches, height, width = masks.shape
        mask = masks[batches - 1]
        mask = (mask > 0.5).astype(np.uint8)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        result = []
        if len(contours) > 0:
            largest_contour = max(contours, key=cv2.contourArea)
            for point in largest_contour:
                x, y = point[0]
                x_scaled = round(x / ratio_width + imageConfigX, 2)
                y_scaled = round(y / ratio_height + imageConfigY, 2)
                result.extend([x_scaled, y_scaled])
        # matplotlib.use('TkAgg')
        # plt.figure(figsize=(10, 10))
        # plt.imshow(image_cv2)
        # self.show_mask(mask, plt.gca())
        # plt.show()

        return result

    def set_image(self, image: np.ndarray):
        input_image = self.transform.apply_image(image)
        input_image_torch = torch.as_tensor(input_image, device=self.device)
        input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
        self.set_torch_image(input_image_torch, image.shape[:2])

    @torch.no_grad()
    def set_torch_image(self, transformed_image: torch.Tensor, original_image_size: Tuple[int, ...]):
        self.reset_image()
        self.original_size = original_image_size
        self.input_size = tuple(transformed_image.shape[-2:])
        input_image = self.model.preprocess(transformed_image)
        self.features = self.model.image_encoder(input_image)
        self.is_image_set = True

    def reset_image(self) -> None:
        self.is_image_set = False
        self.features = None
        self.orig_h = None
        self.orig_w = None
        self.input_h = None
        self.input_w = None

    def _predict(
            self,
            point_coords: Optional[torch.Tensor] = None,
            point_labels: Optional[torch.Tensor] = None,
            box: Optional[torch.Tensor] = None,
            mask_input: Optional[torch.Tensor] = None,
            multi_mask_output: bool = True,
            return_logits: bool = False,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Args:
            point_coords (torch.Tensor or None): point prompts，（BxNx2)，每个point prompt是像素坐标(X,Y)。
            point_labels (torch.Tensor or None): point prompts对应的labels，(BxN)，每个元素是1或0，1代表前景点，0代表背景点。
            box (np.ndarray or None): box prompts，(1x4)，每个box的结构是（X1,Y1,X2,Y2)。
            mask_input (np.ndarray): 低分辨率的mask，一般来自上一次的预测的mask结果，(Bx1xHxW)，H=W=256
            multi_mask_output (bool):
                如果设置为 true，模型将返回三个掩码（mask）。
                对于存在歧义的输入提示（例如单击一次），这种方式通常会比只生成一个掩码的预测效果更好。
                如果只需要单个掩码，可以利用模型预测的质量评分来选择最佳掩码。
                对于没有歧义的提示（例如多个输入提示），将 multimask_output=False 通常可以得到更好的结果。
            return_logits (bool): 如果设置为 true，将返回未经阈值化的掩码 logits，而不是binary mask。
                logits指的是模型输出的原始分数，它保留了更丰富的置信度信息。
                binary_mask指的是0和1，直接表示前景/背景。
        Returns:
            (torch.Tensor): (BxCxHxW)，C代表 masks 的数量，(H,W)是原始图像大小。
            (torch.Tensor): (BxC)，模型针对每个mask的质量预测。
            (torch.Tensor): (BxCxHxW)，一个形状为 B × C × H × W 的数组，其中 C 表示掩码（mask）的数量，且 H = W = 256。
        """
        if not self.is_image_set:
            raise RuntimeError("在进行预测前，需要先调用set_image(...)方法来处理图像。")

        coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
        if point_coords is not None:
            assert (point_labels is not None), "point_labels should not be None."
            point_coords = self.transform.apply_coords(point_coords, self.original_size)
            coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device)
            labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)
            coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]
        if box is not None:
            box = self.transform.apply_boxes(box, self.original_size)
            box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device)
            box_torch = box_torch[None, :]
        if mask_input is not None:
            mask_input_torch = torch.as_tensor(mask_input, dtype=torch.bool, device=self.device)
            mask_input_torch = mask_input_torch[None, :, :, :]
        masks, iou_predictions, low_res_masks = self.predict_torch(
            point_coords=coords_torch,
            point_labels=labels_torch,
            boxes=box_torch,
            mask_input=mask_input_torch,
        )
        masks_np = masks[0].detach().cpu().numpy()
        iou_predictions_np = iou_predictions[0].detach().cpu().numpy()
        low_res_masks_np = low_res_masks[0].detach().cpu().numpy()
        return masks_np, iou_predictions_np, low_res_masks_np

    @torch.no_grad()
    def predict_torch(
            self,
            point_coords: Optional[torch.Tensor] = None,
            point_labels: Optional[torch.Tensor] = None,
            boxes: Optional[torch.Tensor] = None,
            mask_input: Optional[torch.Tensor] = None,
            multi_mask_output: bool = True,
            return_logits: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        基于当前的image和输入prompt来预测masks。输入prompts已被转换成torch.tensor。
        Args:
            point_coords (torch.Tensor or None): point prompts，BxNx2，每个点的像素坐标(X,Y)。
            point_labels (torch.Tensor or None): 对应point prompts的标签，BxN，1代表前景点，0代表背景点。
            boxes (np.ndarray or None): box prompts，Bx4，XYXY格式。
            mask_input (np.ndarray):低分辨率的mask，一般来自上一次预测的结果，Bx1xHxW。
            multi_mask_output (bool):
                如果设置为 true，模型将返回三个掩码（mask）。
                对于存在歧义的输入提示（例如单击一次），这种方式通常会比只生成一个掩码的预测效果更好。
                如果只需要单个掩码，可以利用模型预测的质量评分来选择最佳掩码。
                对于没有歧义的提示（例如多个输入提示），将 multimask_output=False 通常可以得到更好的结果。
            return_logits (bool): 如果设置为 true，将返回未经阈值化的掩码 logits，而不是binary mask。
                logits指的是模型输出的原始分数，它保留了更丰富的置信度信息。
                binary_mask指的是0和1，直接表示前景/背景。
        Returns:
            (torch.Tensor)：输出的 masks，C是mask的数量，H,W是原始图像大小。
            (torch.Tensor)：每个mask的质量评价。BxC。
            (torch.Tensor)：BxCxHxW，其中C是mask的数量。
        """
        if not self.is_image_set:
            raise RuntimeError("在推理前，需要先调用set_image(...)")
        if point_coords is not None:
            points = (point_coords, point_labels)
        else:
            points = None

        sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
            points=points,
            boxes=boxes,
            masks=mask_input
        )

        low_res_masks, iou_predictions = self.model.mask_decoder(
            image_embeddings=self.features,
            image_pe=self.model.prompt_encoder.get_dense_pe(),
            sparse_prompt_embeddings=sparse_embeddings,
            dense_prompt_embeddings=dense_embeddings,
            multi_mask_output=multi_mask_output,
        )
        masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size)
        if not return_logits:
            masks = masks > self.model.mask_threshold
        return masks, iou_predictions, low_res_masks

    def show_mask(self, mask, ax, random_color=False):
        if random_color:
            color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
        else:
            color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
        h, w = mask.shape[-2:]
        mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
        ax.imshow(mask_image)

    def show_points(self, coords, labels, ax, marker_size=375):
        pos_points = coords[labels == 1]
        neg_points = coords[labels == 0]
        ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
                   linewidth=1.25)
        ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
                   linewidth=1.25)
